package cn.cyh.imeans.utils;

import cn.cyh.imeans.enums.ScanEnum;
import org.springframework.util.StringUtils;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.StringJoiner;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class ShellUtil {

    private ShellUtil() {}

    /**
     * 获取脚本执行结果
     * @param dir 脚本目录
     * @param fileName 脚本文件名
     * @return str
     */
    public static String execShell(String dir, String fileName, ScanEnum type, String... params) {
        try {
            String fileFullName = dir + File.separator + fileName;
            File file = new File(fileFullName);
            if (!file.exists()) {
                throw new RuntimeException("shell脚本不存在.");
            }

            List<String> cmd = new ArrayList<>(params.length + 2);
            cmd.add("sh");
            cmd.add(fileName);
            cmd.addAll(Arrays.asList(params));
            System.out.println("执行脚本命令："+ String.join(" ", cmd));

            ProcessBuilder processBuilder = new ProcessBuilder(cmd);
            processBuilder.redirectErrorStream(true);
            Process process = processBuilder.directory(new File(dir)).start();

            StringJoiner sb = new StringJoiner("\n");
            CompletableFuture<Void> zc = CompletableFuture.runAsync(() -> {
                try (BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
                    String line;
                    if (ScanEnum.VIRUS == type) {
//                        log.info("病毒扫描结果：");
                        boolean zw = false;
                        while ((line = reader.readLine()) != null) {
//                            log.info(line);
                            // 取 出现 --- Statistics --- 之后的内容
                            if (line.contains("--- Statistics ---")) {
                                zw = true;
                            }
                            if (zw) {
                                sb.add(line);
                            }
                        }
                    } else if (ScanEnum.VULNERABILITY == type) {
//                        log.info("漏洞扫描结果：");
                        boolean f = true;
                        while ((line = reader.readLine()) != null) {
//                            log.info(line);
                            if(line.contains("Traceback") || line.contains("usage: python")) {
                                f = false;
                            }
                            if(f){
                                sb.add(line);
                            }
                        }
                    } else if (ScanEnum.PORT == type) {
//                        log.info("端口扫描结果：");
                        while ((line = reader.readLine()) != null) {
//                            log.info(line);
                            if (StringUtils.hasLength(line)) {
                                sb.add(line);
                            }
                        }
                    } else if (ScanEnum.ZIP == type) {
//                        log.info("打包脚本执行结果：");
                        while ((line = reader.readLine()) != null) {
//                            log.info(line);
                            if (line.contains("complete")) {
                                sb.add(line);
                            }
                        }
                    } else {
//                        log.info("其他执行结果：");
                        while ((line = reader.readLine()) != null) {
//                            log.info(line);
                            sb.add(line);
                        }
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            });

            // 读取进程异常输出
            CompletableFuture<Void> ex = CompletableFuture.runAsync(() -> {
                try (BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(process.getErrorStream()))) {
                    String line;
                    while ((line = bufferedReader.readLine()) != null) {
//                        log.info(line);
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            });

            boolean exitCode = process.waitFor(6, TimeUnit.HOURS);
            System.out.println("命令执行结果：完成="+ exitCode);
            CompletableFuture.allOf(zc, ex).get(30, TimeUnit.SECONDS);
            return sb.toString();
        } catch (IOException | InterruptedException | ExecutionException | TimeoutException e) {
            e.printStackTrace();
        }
        return null;
    }

}
